#!/usr/bin/env python

#python script to analysis the AST from ssp tool. Xingzhong Stevens

import pygraphviz as pgv
import networkx as nx
import sys
sys.path.append('./plib')
import libsql as sql
import libast as ast
import operator


class AST:
    def __init__(self, name):
        self.name = name.split('/')[-1]
        self.g = nx.DiGraph(name=self.name)


    def draw(self, name= None):
        if name == None:
            name = "ast.png"
        A = nx.to_agraph(self.g)
        A.layout(prog = 'dot')
        print "drawing ast..."
        A.draw(name)

    def work(self):
        self.g =  nx.connected_component_subgraphs(self.g.to_undirected())[0]
        
    def get(self, Sql):
        for x in Sql.ls():
            X = ast.TREE(x)
            self.add_node(X)
        print nx.info(self.g)

    def add_node_1(self, node):
        attr = {}
        if node.name != None:
            attr['label'] = node.name
            attr['style'] = 'filled'
            attr['color'] = 'lightblue'
        elif node.value != None:
            attr['label'] = node.value
            attr['style'] = 'filled'
            attr['color'] = 'lightyellow'
        else:
            attr['label'] = node.type
        addr = node.addr
        self.g.add_node(addr, attr)

    def add_edge_1(self, root, child):
        attr = {}
        if root.num > 0 :
            attr['label'] = root.num
        else :
            attr['label'] = ""
        n1 = root.addr
        self.g.add_edge(n1, child, attr)

    def add_node(self, node):
        self.add_node_1(node)
        if node.op0 != None:
            self.add_edge_1(node, node.op0)
        if node.op1 != None:
            self.add_edge_1(node, node.op1)
        if node.op2 != None:
            self.add_edge_1(node, node.op2)
            


